#include <algorithm>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <list>
#include <map>
#include <set>
#include <sstream>
#include <string>
#include <queue>
#include <vector>
#include <cmath>
#include <boost/program_options.hpp>
#include "MetastableSet.h"
#include "StringSplit.h"
#include "BlockAverage.h"
#include "Interpolation.h"
#include "TrajLoader.h"
#include "Angle.h"

#define MaxLineLength 2048

namespace po = boost::program_options;
using namespace std;

inline int angleIdx (const double & myangle,
		     const double & bin)
{
  return int((myangle + 180) / bin);
}


int main(int argc, char * argv[])
{
  std::string ifile, ofile, jfile, tfile;
  double center, width, bin, sigma, xi, angleDx;
  unsigned nBlock;

  po::options_description desc ("Allow options");
  desc.add_options()
      ("help,h", "print this message")
      ("meta-center,c",  po::value<double > (&center)->default_value (180),   "center of the metastable set")
      ("meta-width,w",  po::value<double > (&width)->default_value (15),   "width of the metastable set")
      ("sigma,s",  po::value<double > (&sigma)->default_value (1.),   "sigma")
      ("xi,x",  po::value<double > (&xi)->default_value (1.),   "xi")
      ("angle-dx,d",  po::value<double > (&angleDx)->default_value (0.001),   "the dx to cal grad angle")
      ("input-j,j", po::value<std::string > (&jfile)->default_value ("table.j.xvg"), "table of J")
      ("bin,r",  po::value<double > (&bin)->default_value (5.),   "the bin size")
      ("num-block,n",  po::value<unsigned > (&nBlock)->default_value (10),   "number of block")
      ("input-angle,f", po::value<std::string > (&ifile)->default_value ("angaver.name"), "file including angle file names")
      ("input-traj,t", po::value<std::string > (&tfile)->default_value ("traj.name"), "file including traj file names")
      ("output,o", po::value<std::string > (&ofile)->default_value ("fmpt.out"), "the output of first mean passage time");
  
      
  po::variables_map vm;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify (vm);
  if (vm.count("help")){
    std::cout << desc<< "\n";
    return 0;
  }

  unsigned nbin = (360. + 0.01 * bin) / bin;
  bin = 360. / double(nbin);
  cout << "# bin size: " << bin << endl;
  cout << "# nbin: " << nbin << endl;
  // record
  vector<vector<double > > record (nbin);
  vector<vector<double > > record_j (nbin);
  cout << "# j table is " << jfile << endl;
  cout << "# angle names is " << ifile << endl;
  cout << "# traj names  is " << tfile << endl;

  double metal, metau;
  metal = center - width;
  if (metal < -180) metal += 360.;
  int tmpn = (metal + 180 + 0.5 * bin) / bin;
  metal = -180 + tmpn * bin;
  metau = center + width;
  if (metau > 180) metau -= 360.;
  tmpn = (metau + 180 + 0.5 * bin) / bin;
  metau = -180 + tmpn * bin;
  cout << "# metastable set: [ " << metal << " , " << metau << " ]" << endl;
  // metastable set
  MetastableSet ms (metal, metau);

  // table of j
  vector<double > xx0, vv0, dd0;
  {
    ifstream fpname0 (jfile.c_str());
    if (!fpname0){
      cerr << "cannot open file " << jfile << endl;
      return 1;
    }
    char nameline [MaxLineLength];
    while (fpname0.getline(nameline, MaxLineLength)){
      if (nameline[0] == '#' || nameline[0] == '@') continue;
      vector<string > words;
      StringOperation::split (string(nameline), words);
      if (words.size() < 3) {
	cerr << "wrong file format of " << jfile << endl;;
	exit (1);
      }
      xx0.push_back (atof(words[0].c_str()));
      vv0.push_back (atof(words[1].c_str()));
      dd0.push_back (atof(words[2].c_str()));
    }
  }
  PiecewisePoly pp;
  Interpolation::piecewiseLinear (xx0, dd0, pp);
  
  // analyze traj
  ifstream fpname (ifile.c_str());
  ifstream fpname_t (tfile.c_str());
  char nameline [MaxLineLength];
  char nameline_t [MaxLineLength];
  if (!fpname){
    std::cerr << "cannot open file " << ifile << std::endl;
    return 1;
  }
  if (!fpname_t){
    std::cerr << "cannot open file " << tfile << std::endl;
    return 1;
  }
  unsigned countFound = 0;
  unsigned countUnFound = 0;
  unsigned countFile = 0;
  unsigned printCount = 0;
  
  while (fpname.getline(nameline, MaxLineLength)){
    if (nameline[0] == '#') continue;
    if (!fpname_t.getline(nameline_t, MaxLineLength)){
      cerr << "inconsistent angle and traj name files... exit" << endl;
      exit(1);
    }
    if (nameline_t[0] == '#') continue;    
    ifstream angname (nameline);
    if (!angname){
      std::cerr << "cannot open file " << nameline << std::endl;
      return 1;
    }
    TrajLoader_xtc traj_l_xtc;
    traj_l_xtc.reinit (nameline_t);
    // cout << "reading file " << nameline << endl;
    if (printCount == 10) {
      printf ("# reading file %s       \r", nameline);
      printCount = 0;
    }
    printCount ++;
    countFile ++;
    vector<double > times;
    vector<double > anglev;
    vector<vector<vector<double > > > c4Coord;
    vector<vector<double > > boxes;
    char valueline [MaxLineLength];
    while (angname.getline(valueline, MaxLineLength)){
      if (valueline[0] == '#' || valueline[0] == '@') continue;
      vector<string > words;
      StringOperation::split (string(valueline), words);
      if (words.size() < 2) {
	cerr << "wrong file format of " << ifile << endl;
	exit (1);
      }
      times .push_back (atof(words[0].c_str()));
      anglev.push_back (atof(words[1].c_str()));
      if (!traj_l_xtc.load()){
	cerr << "inconsistent angle and traj files... exit" << endl;
	exit(1);
      }
      if (fabs(traj_l_xtc.getTime() - times.back()) > 1e-4){
	cerr << "inconsistent times in angle and traj fiels.. exit" << endl;
	exit(1);
      }
      boxes.push_back (traj_l_xtc.getBox());
      vector<vector<double > > tmpc4;
      traj_l_xtc.formCoords (tmpc4);
      c4Coord.push_back(tmpc4);
      AngleCalculator ac (boxes.back());
      if (fabs(anglev.back() -  ac.calAngle (c4Coord.back())) > 1e-1) {
	cerr << "angle v and calculated angle diffs more than 1e-1" << endl;
      }
    }

    double sum_j = 0.;
    if (ms.inSet(anglev[0])) {
      record[angleIdx(anglev[0], bin)].push_back (0.0);
      record_j[angleIdx(anglev[0], bin)].push_back (0.0);
      countFound ++;
    }
    else {
      bool found = true;
      unsigned nextstab = 0;
      do {
	nextstab ++;
	if (nextstab >= anglev.size()){
	  found = false;
	  break;
	}
      } while (! ms.inSet(anglev[nextstab]));
      if (found) {
	countFound ++;
	record [angleIdx(anglev[0], bin)].push_back (-(times[0] - times[nextstab]));
	sum_j = 0.;
	for (unsigned ii = 0; ii < nextstab; ++ii){
	  AngleCalculator ac (boxes[ii]);
	  double tmp = - pp.value_periodic (anglev[ii]) ;
	  double tmpAngle = ac.calAngle (c4Coord[ii]);
	  vector<vector<double > > tmpGrad;
	  ac.calGrad (c4Coord[ii], angleDx, tmpGrad);
	  // vector<vector<double > > yy (4);
	  // vector<double > yy1(3, 0.);
	  // yy[3] = yy[2] = yy[1] = yy[0] = yy1;
	  // yy[0][0] = 0.1;
	  // yy[2][1] = 0.1;
	  // yy[3][1] = yy[3][2] = 0.1;
	  // ac.calGrad (yy, angleDx, tmpGrad);
	  double sumGrad = 0.;
	  for (unsigned kk = 0; kk < tmpGrad.size(); ++kk){
	    for (unsigned ll = 0; ll < 3; ++ll){
	      sumGrad += tmpGrad[kk][ll] * tmpGrad[kk][ll];
	    }
	  }
	  // printf ("time %f, angle %f %f, sumGrad %e, dj is %f\n", times[ii], anglev[ii], tmpAngle, sumGrad, tmp);
	  sum_j += 0.5 * tmp * tmp * sumGrad * (times[1] - times[0]);
	}
	record_j [angleIdx(anglev[0], bin)].push_back (sum_j);
      }
      else {
	countUnFound ++;
      }
    }
  }
  printf ("\n");
  printf ("# read %d files, %d ( %.1f %% ) hit meta, %d ( %.1f %% ) do not hit\n",
	  countFile,
	  countFound, ((double)(countFound))/((double)(countFile)) * 100.,
	  countUnFound, ((double)(countUnFound))/((double)(countFile)) * 100.);

  BlockAverage ba;  
  BlockAverage ba_j;  
  FILE * fout = fopen (ofile.c_str(), "w");
  if (fout == NULL){
    cout << "cannot open file " << ofile<< endl;
    exit (1);
  }
  for (unsigned ii = 0; ii < record.size(); ++ii){
    double myangle = -180 + bin * (ii + 0.5);
    if (record[ii].size() == 0){
      int prevIdx = int(ii);
      do {
	prevIdx --;
      } while (prevIdx >= 0 && record[prevIdx].size() == 0);
      if (prevIdx >= 0){
	fprintf (fout, "%f %e %e\n",
		 myangle,
		 sigma * ba.getAvg() + xi * ba_j.getAvg(), sigma * ba.getAvgError() + xi * ba_j.getAvgError()
		 );
      }
      else{
	fprintf (fout, "%f\n", myangle);
      }
    }
    else if (record[ii].size() < nBlock * 5){
      ba.processData (record[ii], record[ii].size());
      ba_j.processData (record_j[ii], record_j[ii].size());
      fprintf (fout, "%f %e %e    %e %e    %e %e\n",
	       myangle,
	       sigma * ba.getAvg() + xi * ba_j.getAvg(), sigma * ba.getAvgError() + xi * ba_j.getAvgError(),
	       ba.getAvg(), ba.getAvgError(),
	       ba_j.getAvg(), ba_j.getAvgError()
	  );
    }
    else {
      ba.processData (record[ii], nBlock);
      ba_j.processData (record_j[ii], nBlock);
      fprintf (fout, "%f %e %e    %e %e    %e %e\n",
	       myangle,
	       sigma * ba.getAvg() + xi * ba_j.getAvg(), sigma * ba.getAvgError() + xi * ba_j.getAvgError(),
	       ba.getAvg(), ba.getAvgError(),
	       ba_j.getAvg(), ba_j.getAvgError()
	  );
      // fprintf (fout, "%f %e %e\n", myangle, ba.getAvg(), ba.getAvgError());
    }
  }  
  fclose (fout);

  // string ofile1 = ofile + string (".f");
  // fout = fopen (ofile.c_str(), "w");
  // if (fout == NULL){
  //   cout << "cannot open file " << ofile<< endl;
  //   exit (1);
  // }
  // for (unsigned ii = 0; ii < record.size(); ++ii){
  //   double myangle = -180 + bin * (ii + 0.5);
  //   if (record[ii].size() == 0){
  //     fprintf (fout, "%f\n", myangle);
  //   }
  //   else if (record[ii].size() < nBlock * 5){
  //     ba.processData (record[ii], record[ii].size());
  //     ba_j.processData (record_j[ii], record_j[ii].size());
  //     fprintf (fout, "%f %e %e    %e %e    %e %e\n",
  // 	       myangle,
  // 	       sigma * ba.getAvg() + xi * ba_j.getAvg(), sigma * ba.getAvgError() + xi * ba_j.getAvgError(),
  // 	       ba.getAvg(), ba.getAvgError(),
  // 	       ba_j.getAvg(), ba_j.getAvgError()
  // 	  );
  //   }
  //   else {
  //     ba.processData (record[ii], nBlock);
  //     ba_j.processData (record_j[ii], nBlock);
  //     fprintf (fout, "%f %e %e  \n",
  // 	       myangle,
  // 	       sigma * ba.getAvg() + xi * ba_j.getAvg(), sigma * ba.getAvgError() + xi * ba_j.getAvgError(),
  // 	       ba.getAvg(), ba.getAvgError(),
  // 	       ba_j.getAvg(), ba_j.getAvgError()
  // 	  );
  //     // fprintf (fout, "%f %e %e\n", myangle, ba.getAvg(), ba.getAvgError());
  //   }
  // }  
  // fclose (fout);
  
  
  return 0;
}
